package com.tezign.intelligence.mini.dam.aifamilyconsultant.util;

import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import javax.crypto.Cipher;

public class RSAUtils {
	private static final String ALGORITHM = "RSA";
	private static final int KEY_SIZE = 2048; // 密钥长度

	/**
	 * 生成密钥对
	 */
	public static KeyPair generateKeyPair() throws Exception {
		KeyPairGenerator generator = KeyPairGenerator.getInstance(ALGORITHM);
		generator.initialize(KEY_SIZE, new SecureRandom());
		return generator.generateKeyPair();
	}

	/**
	 * 获取公钥字符串
	 */
	public static String getPublicKeyString(PublicKey publicKey) {
		return Base64.getEncoder().encodeToString(publicKey.getEncoded());
	}

	/**
	 * 获取私钥字符串
	 */
	public static String getPrivateKeyString(PrivateKey privateKey) {
		return Base64.getEncoder().encodeToString(privateKey.getEncoded());
	}

	/**
	 * 从字符串还原公钥
	 */
	public static PublicKey getPublicKey(String publicKeyStr) throws Exception {
		byte[] keyBytes = Base64.getDecoder().decode(publicKeyStr);
		X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
		KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
		return keyFactory.generatePublic(keySpec);
	}

	/**
	 * 从字符串还原私钥
	 */
	public static PrivateKey getPrivateKey(String privateKeyStr) throws Exception {
		byte[] keyBytes = Base64.getDecoder().decode(privateKeyStr);
		PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
		KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
		return keyFactory.generatePrivate(keySpec);
	}

	/**
	 * 公钥加密
	 */
	public static String encrypt(String content, PublicKey publicKey) throws Exception {
		Cipher cipher = Cipher.getInstance(ALGORITHM);
		cipher.init(Cipher.ENCRYPT_MODE, publicKey);
		byte[] encrypted = cipher.doFinal(content.getBytes());
		return Base64.getEncoder().encodeToString(encrypted);
	}

	/**
	 * 私钥解密
	 */
	public static String decrypt(String encrypted, PrivateKey privateKey) throws Exception {
		Cipher cipher = Cipher.getInstance(ALGORITHM);
		cipher.init(Cipher.DECRYPT_MODE, privateKey);
		byte[] decrypted = cipher.doFinal(Base64.getDecoder().decode(encrypted));
		return new String(decrypted);
	}

	public static void main(String[] args) throws Exception {
		// 生成密钥对
		KeyPair keyPair = generateKeyPair();
		PublicKey publicKey = keyPair.getPublic();
		PrivateKey privateKey = keyPair.getPrivate();

		// 获取密钥字符串
		String publicKeyString = getPublicKeyString(publicKey);
		String privateKeyString = getPrivateKeyString(privateKey);

		System.out.println("Public Key:");
		System.out.println(publicKeyString);
		System.out.println("\nPrivate Key:");
		System.out.println(privateKeyString);

		// 测试加解密
		String originalText = "Hello, RSA!";
		System.out.println("\nOriginal: " + originalText);

		// 加密
		String encrypted = encrypt(originalText, publicKey);
		System.out.println("Encrypted: " + encrypted);

		// 解密
		String decrypted = decrypt(encrypted, privateKey);
		System.out.println("Decrypted: " + decrypted);
	}
}
